"use client";

import { useEffect, useRef, useState, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';

interface WebSocketMessage {
  id: string;
  role: 'USER' | 'ASSISTANT';
  content: string;
  metadata?: any;
  createdAt: string;
  userId?: string;
}

interface UseWebSocketOptions {
  sessionId: string;
  accessToken: string;
  onMessage?: (message: WebSocketMessage) => void;
  onError?: (error: string) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
}

export function useWebSocket({
  sessionId,
  accessToken,
  onMessage,
  onError,
  onConnect,
  onDisconnect,
}: UseWebSocketOptions) {
  const [isConnected, setIsConnected] = useState(false);
  const [connectionError, setConnectionError] = useState<string | null>(null);
  const socketRef = useRef<Socket | null>(null);

  const connect = useCallback(() => {
    // If already connected, disconnect first
    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
    }

    const socket = io(`${process.env.NEXT_PUBLIC_WS_URL || 'ws://localhost:3001'}/chat`, {
      transports: ['websocket'],
      timeout: 10000,
    });

    socket.on('connect', () => {
      setIsConnected(true);
      setConnectionError(null);
      
      // Send authentication info
      if (accessToken && sessionId) {
        socket.emit('authenticate', { accessToken, sessionId }, (response: any) => {
          if (response.success) {
            onConnect?.();
          } else {
            setConnectionError(response.error || 'Authentication failed');
            onError?.(response.error || 'Authentication failed');
          }
        });
      }
    });

    socket.on('disconnect', () => {
      setIsConnected(false);
      onDisconnect?.();
    });

    socket.on('connect_error', (error) => {
      console.error('WebSocket connection error:', error);
      setConnectionError(error.message || 'Connection failed');
      onError?.(error.message || 'Connection failed');
    });

    socket.on('message_received', (message: WebSocketMessage) => {
      onMessage?.(message);
    });

    socket.on('tts_chunk', (data: { bytes: ArrayBuffer; sessionId: string }) => {
      // Voice data processing - to be implemented
    });

    socket.on('utterance_ended', (data: { sessionId: string; timestamp: string }) => {
      // Voice end event handling - to be implemented
    });

    socketRef.current = socket;
  }, [sessionId, accessToken, onMessage, onError, onConnect, onDisconnect]);

  const disconnect = useCallback(() => {
    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
      setIsConnected(false);
    }
  }, []);

  const sendMessage = useCallback((content: string, metadata?: any) => {
    if (!socketRef.current?.connected) {
      throw new Error('WebSocket not connected');
    }

    return new Promise((resolve, reject) => {
      socketRef.current!.emit('send_message', { content, metadata }, (response: any) => {
        if (response.success) {
          resolve(response);
        } else {
          reject(new Error(response.error || 'Failed to send message'));
        }
      });
    });
  }, []);

  const sendVoiceChunk = useCallback((bytes: ArrayBuffer, meta?: any) => {
    if (!socketRef.current?.connected) {
      throw new Error('WebSocket not connected');
    }

    socketRef.current.emit('voice_chunk', {
      sessionId,
      bytes,
      meta,
    });
  }, [sessionId]);

  const endUtterance = useCallback(() => {
    if (!socketRef.current?.connected) {
      throw new Error('WebSocket not connected');
    }

    socketRef.current.emit('end_of_utterance', { sessionId });
  }, [sessionId]);

  const getOnlineCount = useCallback(() => {
    if (!socketRef.current?.connected) {
      throw new Error('WebSocket not connected');
    }

    return new Promise((resolve, reject) => {
      socketRef.current!.emit('get_online_count', {}, (response: any) => {
        if (response.success) {
          resolve(response.count);
        } else {
          reject(new Error(response.error || 'Failed to get online count'));
        }
      });
    });
  }, []);

  useEffect(() => {
    if (sessionId && accessToken) {
      connect();
    }

    return () => {
      disconnect();
    };
  }, [sessionId, accessToken]); // Depends on sessionId and accessToken

  return {
    isConnected,
    connectionError,
    sendMessage,
    sendVoiceChunk,
    endUtterance,
    getOnlineCount,
    reconnect: connect,
  };
}
